{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quickstart"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This quickstart is intended for developers who are ready to dive into the code and see an example of how to train and use 🤗 SetFit models. We recommend starting with this quickstart, and then proceeding to the [tutorials](https://huggingface.co/docs/setfit/main/en/./tutorials/overview) or [how-to guides](https://huggingface.co/docs/setfit/main/en/./how_to/overview) for additional material. Additionally, the [conceptual guides](https://huggingface.co/docs/setfit/main/en/./conceptual_guides/setfit) help explain exactly how SetFit works.\n",
    "\n",
    "Start by installing 🤗 SetFit:\n",
    "\n",
    "```bash\n",
    "pip install setfit\n",
    "```\n",
    "\n",
    "If you have a CUDA-capable graphics card, then it is recommended to [install `torch` with CUDA support](https://pytorch.org/get-started/locally/) to train and performing inference much more quickly:\n",
    "\n",
    "```bash\n",
    "pip install torch --index-url https://download.pytorch.org/whl/cu118\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SetFit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SetFit is an efficient framework to train low-latency text classification models using little training data. In this Quickstart, you'll learn how to train a SetFit model, how to perform inference with it, and how to save it to the Hugging Face Hub."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, you'll load a [Sentence Transformer model](https://huggingface.co/models?library=sentence-transformers) and further finetune it for classifying movie reviews as positive or negative. To train a model, we will need to prepare the following three: 1) a **model**, 2) a **dataset**, and 3) **training arguments**.\n",
    "\n",
    "**1**. Initialize a SetFit model using a Sentence Transformer model of our choice. Consider using the [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard) to guide your decision on which Sentence Transformer model to choose. We will use [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5), a small but performant model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from setfit import SetFitModel\n",
    "\n",
    "model = SetFitModel.from_pretrained(\"BAAI/bge-small-en-v1.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2a**. Next, load both the \"train\" and \"test\" splits of the [SetFit/sst2](https://huggingface.co/datasets/sst2) dataset. Note that the dataset has `\"text\"` and `\"label\"` columns: this is exactly the format that 🤗 SetFit expects. If your dataset has different columns, then you can use the column_mapping argument of the [Trainer](https://huggingface.co/docs/setfit/main/en/reference/trainer#setfit.Trainer) in step 4 to map the column names to `\"text\"` and `\"label\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text', 'label', 'label_text'],\n",
       "        num_rows: 6920\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['text', 'label', 'label_text'],\n",
       "        num_rows: 1821\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['text', 'label', 'label_text'],\n",
       "        num_rows: 872\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"SetFit/sst2\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2b**. In real world scenarios it is very uncommon to have ~7.000 high quality labeled training samples, so we will heavily shrink the training dataset to give a better idea of how 🤗 SetFit would work in real settings. To be specific, the `sample_dataset` function will sample only 8 samples for each class. The testing set is left unaffected for better evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['text', 'label', 'label_text'],\n",
       "    num_rows: 16\n",
       "})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from setfit import sample_dataset\n",
    "\n",
    "train_dataset = sample_dataset(dataset[\"train\"], label_column=\"label\", num_samples=8)\n",
    "train_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['text', 'label', 'label_text'],\n",
       "    num_rows: 1821\n",
       "})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_dataset = dataset[\"test\"]\n",
    "test_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2c**. We can apply the labels from the dataset on the model, so the predictions output readable classes. You can also provide the labels directly to `SetFitModel.from_pretrained()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.labels = [\"negative\", \"positive\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**3**. Prepare the [TrainingArguments](https://huggingface.co/docs/setfit/main/en/reference/trainer#setfit.TrainingArguments) for training. Note that training with 🤗 SetFit consists of two phases behind the scenes: **finetuning embeddings** and **training a classification head**. As a result, some of the training arguments can be tuples, where the two values are used for each of the two phases, respectively.\n",
    "\n",
    "The `num_epochs` and `max_steps` arguments are frequently used to increase and decrease the number of total training steps. Consider that with SetFit, better performance is reached with **more data, not more training**! Don't be afraid to train for less than 1 epoch if you have a lot of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from setfit import TrainingArguments\n",
    "\n",
    "args = TrainingArguments(\n",
    "    batch_size=32,\n",
    "    num_epochs=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**4**. Initialize the [Trainer](https://huggingface.co/docs/setfit/main/en/reference/trainer#setfit.Trainer) and perform training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from setfit import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=train_dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "***** Running training *****\n",
       "  Num examples = 5\n",
       "  Num epochs = 10\n",
       "  Total optimization steps = 50\n",
       "  Total train batch size = 32\n",
       "{'embedding_loss': 0.2077, 'learning_rate': 4.000000000000001e-06, 'epoch': 0.2}                                                                                                                \n",
       "{'embedding_loss': 0.0097, 'learning_rate': 0.0, 'epoch': 10.0}                                                                                                                                 \n",
       "{'train_runtime': 14.705, 'train_samples_per_second': 108.807, 'train_steps_per_second': 3.4, 'epoch': 10.0}\n",
       "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:08<00:00,  5.70it/s]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**5**. Perform evaluation using the provided testing dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "***** Running evaluation *****\n",
       "{'accuracy': 0.8511806699615596}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feel free to experiment with increasing the number of samples per class to observe the improvements in accuracy. As a challenge, you can play with the samples per class, learning rate, number of epochs, maximum number of steps, and the base Sentence Transformer model to try and improve the accuracy over 90% using very little data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving a 🤗 SetFit model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After training, you can save a 🤗 SetFit model to your local filesystem or to the Hugging Face Hub. Save a model to a local directory using `SetFitModel.save_pretrained()` by providing a `save_directory`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"setfit-bge-small-v1.5-sst2-8-shot\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, push a model to the Hugging Face Hub using `SetFitModel.push_to_hub()` by providing a `repo_id`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.push_to_hub(\"tomaarsen/setfit-bge-small-v1.5-sst2-8-shot\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading a 🤗 SetFit model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A 🤗 SetFit model can be loaded using `SetFitModel.from_pretrained()` by providing 1) a `repo_id` from the Hugging Face Hub or 2) a path to a local directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SetFitModel.from_pretrained(\"tomaarsen/setfit-bge-small-v1.5-sst2-8-shot\") # Load from the Hugging Face Hub\n",
    "\n",
    "model = SetFitModel.from_pretrained(\"setfit-bge-small-v1.5-sst2-8-shot\") # Load from a local directory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once a 🤗 SetFit model has been trained, then it can be used for inference to classify reviews using [SetFitModel.predict()](https://huggingface.co/docs/setfit/main/en/reference/main#setfit.SetFitModel.predict) or [SetFitModel.__call__()](https://huggingface.co/docs/setfit/main/en/reference/main#setfit.SetFitModel.__call__):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['positive' 'negative' 'negative']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds = model.predict([\n",
    "    \"It's a charming and often affecting journey.\",\n",
    "    \"It's slow -- very, very slow.\",\n",
    "    \"A sometimes tedious film.\",\n",
    "])\n",
    "preds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These predictions rely on the `model.labels`. If not set, it will return predictions in the format that was used during training, e.g. `tensor([1, 0, 0])`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's next?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You've completed the 🤗 SetFit quickstart! You can train, save, load and perform inference with 🤗 SetFit models!\n",
    "\n",
    "For your next steps, take a look at our [How-to guides](https://huggingface.co/docs/setfit/main/en/./how_to/overview) and learn how to do more specific things like hyperparameter search, knowledge distillation, or zero-shot text classification. If you're interested in learning more about how 🤗 SetFit works, grab a cup of coffee and read our [Conceptual Guides](https://huggingface.co/docs/setfit/main/en/./conceptual_guides/setfit)!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End-to-end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This snippet shows the entire quickstart in an end-to-end example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from setfit import SetFitModel, Trainer, TrainingArguments, sample_dataset\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Initializing a new SetFit model\n",
    "model = SetFitModel.from_pretrained(\"BAAI/bge-small-en-v1.5\", labels=[\"negative\", \"positive\"])\n",
    "\n",
    "# Preparing the dataset\n",
    "dataset = load_dataset(\"SetFit/sst2\")\n",
    "train_dataset = sample_dataset(dataset[\"train\"], label_column=\"label\", num_samples=8)\n",
    "test_dataset = dataset[\"test\"]\n",
    "\n",
    "# Preparing the training arguments\n",
    "args = TrainingArguments(\n",
    "    batch_size=32,\n",
    "    num_epochs=10,\n",
    ")\n",
    "\n",
    "# Preparing the trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=train_dataset,\n",
    ")\n",
    "trainer.train()\n",
    "\n",
    "# Evaluating\n",
    "metrics = trainer.evaluate(test_dataset)\n",
    "print(metrics)\n",
    "# => {'accuracy': 0.8511806699615596}\n",
    "\n",
    "# Saving the trained model\n",
    "model.save_pretrained(\"setfit-bge-small-v1.5-sst2-8-shot\")\n",
    "# or\n",
    "model.push_to_hub(\"tomaarsen/setfit-bge-small-v1.5-sst2-8-shot\")\n",
    "\n",
    "# Loading a trained model\n",
    "model = SetFitModel.from_pretrained(\"tomaarsen/setfit-bge-small-v1.5-sst2-8-shot\") # Load from the Hugging Face Hub\n",
    "# or\n",
    "model = SetFitModel.from_pretrained(\"setfit-bge-small-v1.5-sst2-8-shot\") # Load from a local directory\n",
    "\n",
    "# Performing inference\n",
    "preds = model.predict([\n",
    "    \"It's a charming and often affecting journey.\",\n",
    "    \"It's slow -- very, very slow.\",\n",
    "    \"A sometimes tedious film.\",\n",
    "])\n",
    "print(preds)\n",
    "# => [\"positive\", \"negative\", \"negative\"]"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
